//! label dictionary

use std::collections::HashMap;

use anyhow::{Result, bail};

/// label dictionary to store the labels' addresses
#[derive(Debug)]
pub(super) struct LabelDictionary {
    labels: HashMap<String, u16>,
    variable_counter: u16,
}

impl LabelDictionary {
    pub fn new() -> Self {
        LabelDictionary {
            labels: HashMap::from([
                ("R0".into(), 0),
                ("R1".into(), 1),
                ("R2".into(), 2),
                ("R3".into(), 3),
                ("R4".into(), 4),
                ("R5".into(), 5),
                ("R6".into(), 6),
                ("R7".into(), 7),
                ("R8".into(), 8),
                ("R9".into(), 9),
                ("R10".into(), 10),
                ("R11".into(), 11),
                ("R12".into(), 12),
                ("R13".into(), 13),
                ("R14".into(), 14),
                ("R15".into(), 15),
                ("SCREEN".into(), 16384),
                ("KBD".into(), 24576),
                ("SP".into(), 0),
                ("LCL".into(), 1),
                ("ARG".into(), 2),
                ("THIS".into(), 3),
                ("THAT".into(), 4),
            ]),
            variable_counter: 16,
        }
    }
    /// Add a label to the dictionary if the label is not already present
    pub fn add(&mut self, code: &str, address: u16) -> Result<()> {
        if !code.starts_with('(') || !code.ends_with(')') {
            bail!("Invaid label instruction {code}")
        }
        let label = &code[1..code.len() - 1];
        if !self.labels.contains_key(label) {
            self.labels.insert(label.to_string(), address);
        }
        Ok(())
    }
    /// get the address of a label
    /// if it not exists, set it's address to self.variable_counter
    /// and add it to the dictionary
    pub fn get(&mut self, label: &str) -> u16 {
        self.labels
            .entry(label.into())
            .or_insert_with_key(|_| {
                self.variable_counter += 1;
                self.variable_counter - 1
            })
            .to_owned()
    }
}
